# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: junglekim

import warnings
warnings.filterwarnings("ignore")
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.ticker as plticker
from matplotlib.collections import LineCollection
from matplotlib.cm import ScalarMappable
from matplotlib.font_manager import FontProperties

TYPE_MAP = {
  'ORDER_ACCEPTED': 1,
  'ORDER_FILLED': 2,
  'CANCEL_CONFIRMED': 3
}

DEPTH_MIN_BP = -350
DEPTH_MAX_BP = 350
INTERVAL = 50

def plot_lp_plot(filename_prefix,
                  start_time,
                  end_time,
                  oedf,
                  bbodf,
                  plot_dir="lp_plot"):
  os.makedirs(plot_dir, exist_ok=True)
  filename = os.path.join(plot_dir, "lpplt_%s.png" % filename_prefix)
  start_ts, end_ts = pd.DatetimeIndex([start_time, end_time]).astype(int)

  print("\nProcessing %s ..." % filename_prefix)
  arr, data, midp = gen_data(oedf, bbodf, start_ts, end_ts)
  if (arr is None) & (data is None) & (midp is None):
    return pd.Series(), 0

  plt.rcParams['figure.figsize'] = 15, 17
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 9
  plt.rcParams['xtick.labelsize'] = 9
  plt.rcParams['ytick.labelsize'] = 9
  plt.rcParams['axes.facecolor'] = 'black'
  plt.rcParams['figure.facecolor'] = 'black'
  plt.rcParams['savefig.facecolor']='black'
  plt.rcParams['xtick.color'] = 'white'
  plt.rcParams['ytick.color'] = 'white'

  plt.subplot(10, 1, (1, 5))
  plt.title(f"{filename_prefix}", c='white')
  plot_order_plot(arr, midp, start_ts, end_ts)

  plt.subplot(10, 1, (6, 8))
  plot_liquidity(data, start_ts, end_ts)
  
  plt.subplot(10, 1, (9, 10))
  tbarr, intervals = plot_depth_tb(data=data)
  plt.tight_layout()
  plt.savefig(filename, dpi=250)
  
  plt.close()
  plt.rcParams.update(plt.rcParamsDefault)
  return tbarr, intervals

def gen_data(oedf, bbodf, startts, endts):
  def gen(x):
    if (len(x) == 1) | (1 not in x['type'].values):
      return None
    return np.array([i for i in zip(x['timestamp'], x['timestamp'].shift(-1), x['order_price'], x['value'])][:-1])

  def spread_value(cols, idxs, arr):
    varr = np.zeros((len(idxs), len(cols)))
    for t1, t2, price, value in arr:
      varr[np.where((idxs > t1) & (idxs < t2)), np.where(cols == price)] += value
    return varr

  cpdf = oedf.copy()
  cpbbodf = bbodf.copy()
  cpdf = cpdf.iloc[1:-1]
  
  cpdf = cpdf.loc[(startts < cpdf['timestamp']) & (cpdf['timestamp'] < endts)]
  cpbbodf = cpbbodf.loc[(startts < cpbbodf['timestamp']) & (cpbbodf['timestamp'] < endts)]

  cpdf = cpdf[cpdf['type'].isin(TYPE_MAP.keys())]
  cpdf['type'] = cpdf['type'].map(TYPE_MAP)
  cpdf['fill_qty'][cpdf['type'] % 2 == 1] = np.nan
  
  midp = cpbbodf[['timestamp', 'mid_price']]
  midp['timestamp'] = pd.DatetimeIndex(midp['timestamp'])
  midp = midp.set_index('timestamp').dropna().resample('5s').last().ffill()
  midp.index = midp.index.astype(int)

  cpdf = cpdf.set_index(['order_id', 'type', 'timestamp']).sort_index(level=['order_id', 'type', 'timestamp'],
                                                                      ascending=[True, True, True])[['order_price', 'order_qty', 'fill_qty']]
  cpdf['value'] = cpdf['order_price'] * (-cpdf['fill_qty']).fillna(cpdf['order_qty'])
  cpdf = cpdf[['order_price', 'value']]
  cpdf = pd.concat([cpdf.groupby('order_id').agg({'value':'cumsum'}), cpdf['order_price']], axis=1)
  arr = cpdf.reset_index().groupby('order_id').apply(gen).values
  arr = [elem for elem in arr if elem is not None]
  if len(arr) == 0:
    return None, None, None
  arr = np.vstack(np.array(arr, dtype=object)).astype(float)
  cols = cpdf['order_price'].dropna().unique()
  idxs = midp.index.to_numpy()
  midps = midp['mid_price'].values

  varr = spread_value(cols, idxs, arr)

  bins = np.linspace(DEPTH_MIN_BP, DEPTH_MAX_BP, int((DEPTH_MAX_BP - DEPTH_MIN_BP) / INTERVAL + 1))
  labels = bins[bins != 0]
  rarr = np.array([])
  rarr = np.stack([np.digitize((price / midps - 1) * 1e4, bins) for price in cols]).T
  rarr = np.where(rarr == 0, 1, rarr)
  rarr = np.where(rarr == len(bins), len(bins) - 1, rarr)

  farr = np.array([])
  farr = np.stack([np.where(rarr == d, varr, 0).sum(axis=1) for d in np.arange(len(labels)) + 1]).T
  fdf = pd.DataFrame(farr, index=idxs, columns=labels)

  return arr[~np.any(arr<=0, axis=1)], fdf, midp

def plot_order_plot(arr, midp, start_ts, end_ts):
  norm = colors.LogNorm(arr[:,-1].min(), arr[:,-1].max())
  lc = LineCollection([line for line in zip(np.concatenate((arr[:,0][:,np.newaxis],
                                                            arr[:,2][:,np.newaxis]),
                                                           axis=1), 
                                            np.concatenate((arr[:,1][:,np.newaxis],
                                                            arr[:,2][:,np.newaxis]),
                                                           axis=1))],
                        cmap='jet',
                        norm=norm,
                        linewidths=2)
  lc.set_array(arr[:,-1])
  line = plt.gca().add_collection(lc)
  plt.plot(midp.index,
           midp['mid_price'],
           c='white')
  plt.plot(midp.index,
           midp['mid_price'] * (1 + 100 * 1e-4),
           c='white',
           linewidth=0.2)
  plt.plot(midp.index,
           midp['mid_price'] / (1 + 100 * 1e-4),
           c='white',
           linewidth=0.2)
  loc = plticker.MultipleLocator(base=3.6e+12) # 4hours
  plt.gca().xaxis.set_major_locator(loc)
  plt.gca().set_xticklabels([pd.to_datetime(x).strftime('%m-%d %H') for x in plt.gca().get_xticks()])
  plt.gca().set_ylabel('prices', c='white', fontsize=14)
  plt.autoscale()
  cbar = plt.colorbar(line, ax=plt.gca(), shrink=0.5, fraction=0.05)
  cbar.set_label('value (USD)', c='white', fontsize=12)
  cbar.outline.set_edgecolor('black')
  [s.set_visible(False) for s in plt.gca().spines.values()]
  plt.xlim(start_ts, end_ts)
  plt.grid(c='gray')

def plot_liquidity(data, start_ts, end_ts):
  df = data.copy()
  df[df.columns[df.columns<0]] *= -1

  cmaplist = ['#8A307F', '#CC313D', '#EC449B', '#6883BC', '#99F443', '#79A7D3', '#F7C5CC']

  plt.stackplot(df.index, df[df.columns[df.columns<0]].iloc[:,::-1].values.T, colors=cmaplist)
  plt.stackplot(df.index, df[df.columns[df.columns>0]].values.T, colors=cmaplist)
  loc = plticker.MultipleLocator(base=3.6e+12) # 4hours
  plt.gca().xaxis.set_major_locator(loc)
  plt.gca().set_xticklabels([pd.to_datetime(x).strftime('%m-%d %H') for x in plt.gca().get_xticks()], color='white')
  plt.gca().set_ylabel('liquidity (USD)', c='white', fontsize=14)
  
  cmap = colors.LinearSegmentedColormap.from_list('depth', cmaplist[:int((DEPTH_MAX_BP - DEPTH_MIN_BP) / INTERVAL)][::-1] + cmaplist[:int((DEPTH_MAX_BP - DEPTH_MIN_BP) / INTERVAL)])
  bounds = np.arange(DEPTH_MIN_BP, DEPTH_MAX_BP + INTERVAL, INTERVAL)
  norm = colors.BoundaryNorm(bounds, cmap.N)
  cbar = plt.colorbar(ScalarMappable(norm=norm, cmap=cmap), ax=plt.gca(), shrink=0.8, ticks=bounds[bounds!=0], fraction=0.05)
  cbar.ax.set_yticklabels([str(i) + 'bp' if i < 0 else '+' + str(i) for i in bounds[bounds!=0]])
  cbar.set_label('edge', c='white', fontsize=12)
  cbar.outline.set_edgecolor('black')
  [s.set_visible(False) for s in plt.gca().spines.values()]
  plt.xlim(start_ts, end_ts)
  plt.gca().grid(c='gray')

def plot_depth_tb(data=None, sumdf=None, interval_num=None):
  assert (data is not None) or ((sumdf is not None) and (interval_num is not None))
  if sumdf is None:
    df = data.copy()
    interval_num = len(df)
    sumdf = df.sum()
  tbarr = (sumdf/interval_num).reset_index().applymap('{:,.0f}'.format).values[::-1]
  tb = plt.table(np.concatenate(np.split(tbarr, 2), axis=1),
                  colWidths=[0.4] * len(tbarr[0]) * 2,
                  loc='center',
                  cellLoc='center')
  for key, cell in tb.get_celld().items():
    cell.PAD = 0.0001
    cell.set(facecolor='black', linewidth=0)
    cell.get_text().set_color('white')
    cell.set_text_props(fontproperties=FontProperties(weight='bold', size=20))
  tb.scale(0.25, 2)
  [s.set_visible(False) for s in plt.gca().spines.values()]
  plt.gca().set_axis_off()

  return sumdf, interval_num

def plot_total_depth_tb(filename_prefix,
                        out_dir,
                        sumdf,
                        interval_num):
  os.makedirs(out_dir, exist_ok=True)
  filename = os.path.join(out_dir, 'depth_%s.png' % filename_prefix)

  print("\nProcessing %s ..." % filename_prefix)

  plt.rcParams['figure.figsize'] = 8, 8
  plt.rcParams['font.size'] = 6
  plt.rcParams['legend.fontsize'] = 6
  plt.rcParams['xtick.labelsize'] = 6
  plt.rcParams['ytick.labelsize'] = 6
  plt.rcParams['axes.facecolor'] = 'black'
  plt.rcParams['figure.facecolor'] = 'black'
  plt.rcParams['savefig.facecolor']='black'
  plt.rcParams['xtick.color'] = 'white'
  plt.rcParams['ytick.color'] = 'white'
  plt.subplot(1, 1, 1)
  plt.title(f"{filename_prefix}", c='white')
  plot_depth_tb(sumdf=sumdf, interval_num=interval_num)
  plt.savefig(filename, dpi=250)
  plt.close()
  plt.rcParams.update(plt.rcParamsDefault)

